import os
import subprocess
import socket

# 定义需要检查的服务和端口
SERVICES = {
    "Print Spooler": {
        "name": "Spooler",
        "impact": "无法提交打印任务，打印任务会卡在队列中无法执行。"
    },
    "Server": {
        "name": "LanmanWorkstation",
        "impact": "无法共享文件和打印机，其他设备无法访问共享打印机。"
    },
    "Workstation": {
        "name": "LanmanWorkstation",
        "impact": "计算机无法作为工作站连接到其他网络资源，包括打印机。"
    },
    "TCP/IP NetBIOS Helper": {
        "name": "lmhosts",
        "impact": "可能导致NetBIOS通信问题，影响打印机的发现和连接。"
    },
    "Remote Procedure Call (RPC)": {
        "name": "RpcSs",
        "impact": "RPC服务未运行时，会导致'RPC服务器不可用'的错误，打印机无法接收打印命令。"
    },
    "DNS Client": {
        "name": "Dnscache",
        "impact": "可能导致域名解析失败，影响打印机的网络发现。"
    },
    "Windows Management Instrumentation": {
        "name": "Winmgmt",
        "impact": "可能导致系统管理信息无法正确获取，影响打印机状态的监控。"
    },
    "Function Discovery Resource Publication": {
        "name": "fdPHost",
        "impact": "可能导致设备和服务信息无法正确发布，影响打印机的自动发现。"
    }
}

PORTS = {
    139: {
        "description": "NetBIOS Session Service",
    },
    445: {
        "description": "SMB",
    },
    9100: {
        "description": "Direct TCP/IP Printing",
    },
    515: {
        "description": "LPD",
    },
    135: {
        "description": "RPC",
    },
    137: {
        "description": "NetBIOS Name Service",
    },
    138: {
        "description": "NetBIOS Datagram Service",
    },
    631: {
        "description": "IPP",
    }
}

def check_service_status(service_name):
    """检查服务状态"""
    try:
        result = subprocess.run(
            ["sc", "query", service_name],
            capture_output=True,
            text=True,
            check=True
        )
        if "RUNNING" in result.stdout:
            return True
        else:
            return False
    except subprocess.CalledProcessError:
        return False

def start_service(service_name):
    """启动服务"""
    try:
        subprocess.run(
            ["sc", "start", service_name],
            check=True
        )
        print(f"服务 {service_name} 已启动。")
    except subprocess.CalledProcessError as e:
        print(f"启动服务 {service_name} 失败: {e}")

def delete_firewall_rule(port):
    """删除指定端口的防火墙规则"""
    try:
        # 删除入站规则
        subprocess.run(
            ["netsh", "advfirewall", "firewall", "delete", "rule", "name=Allow_Port_{}_In".format(port)],
            check=True
        )
        print(f"已删除入站防火墙规则，端口 {port}。")

        # 删除出站规则
        subprocess.run(
            ["netsh", "advfirewall", "firewall", "delete", "rule", "name=Allow_Port_{}_Out".format(port)],
            check=True
        )
        print(f"已删除出站防火墙规则，端口 {port}。")
    except subprocess.CalledProcessError as e:
        print(f"删除防火墙规则失败: {e}")

def add_firewall_rule(port, description):
    """为指定端口添加防火墙规则"""
    try:
        # 添加入站规则
        subprocess.run(
            ["netsh", "advfirewall", "firewall", "add", "rule", "name=Allow_Port_{}_In".format(port),
             "dir=in", "action=allow", "protocol=TCP", "localport={}".format(port)],
            check=True
        )
        print(f"已添加入站防火墙规则，允许端口 {port} ({description}) 的流量。")

        # 添加出站规则
        subprocess.run(
            ["netsh", "advfirewall", "firewall", "add", "rule", "name=Allow_Port_{}_Out".format(port),
             "dir=out", "action=allow", "protocol=TCP", "localport={}".format(port)],
            check=True
        )
        print(f"已添加出站防火墙规则，允许端口 {port} ({description}) 的流量。")
    except subprocess.CalledProcessError as e:
        print(f"添加防火墙规则失败: {e}")

def check_port_status(port):
    """检查端口是否被占用"""
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        try:
            s.bind(("0.0.0.0", port))
            return False  # 端口未被占用
        except socket.error:
            return True  # 端口已被占用

def main():
    print("开始检查服务状态...")
    for service_display_name, service_info in SERVICES.items():
        service_name = service_info["name"]
        if not check_service_status(service_name):
            print(f"服务 {service_display_name} 未运行，可能导致以下问题：{service_info['impact']}")
            print(f"尝试启动服务 {service_display_name}...")
            start_service(service_name)
        else:
            print(f"服务 {service_display_name} 正在运行。")

    print("\n开始删除旧的防火墙规则...")
    for port in PORTS.keys():
        delete_firewall_rule(port)

    print("\n开始添加新的防火墙规则...")
    for port, port_info in PORTS.items():
        add_firewall_rule(port, port_info["description"])

    print("\n开始检查端口状态...")
    for port, port_info in PORTS.items():
        if check_port_status(port):
            print(f"端口 {port} ({port_info['description']}) 被占用")
        else:
            print(f"端口 {port} ({port_info['description']}) 未被占用。")

if __name__ == "__main__":
    main()
